import os
import asyncio
import sys
from typing import Dict, Any, List, Optional
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

# 强制使用httpx作为HTTP客户端
os.environ['DASHSCOPE_HTTP_CLIENT'] = 'httpx'

# 修复OpenSSL Uplink错误
os.environ['OPENSSL_CONF'] = ''
os.environ['SSLKEYLOGFILE'] = ''

import dashscope
from dashscope import Generation
from app.settings import settings
from app.bailian.common import llm


class RealBrowserClient:
    """真实浏览器客户端，使用Playwright MCP服务器"""
    
    def __init__(self):
        self.initialized = False
        self.session = None
        self.tools = None
        self.agent = None
    
    async def initialize(self):
        """初始化真实浏览器客户端"""
        try:
            # 临时修改Python路径，确保导入正确的mcp包
            original_path = sys.path.copy()
            # 移除可能导致冲突的路径
            sys.path = [p for p in sys.path if 'app\\mcp' not in p and 'app/mcp' not in p]
            
            # 动态导入必要的模块
            from mcp import StdioServerParameters, stdio_client, ClientSession
            from langchain_mcp_adapters.tools import load_mcp_tools
            from langgraph.prebuilt import create_react_agent
            
            # 使用npx启动远程MCP服务器
            server_params = StdioServerParameters(
                command="npx",
                args=["-y", "@executeautomation/playwright-mcp-server"]
            )
            
            async with stdio_client(server_params) as (read, write):
                async with ClientSession(read, write) as session:
                    self.session = session
                    await session.initialize()
                    # 获取MCP工具
                    self.tools = await load_mcp_tools(session)
                    print("可用的浏览器工具:", [tool.name for tool in self.tools])
                    
                    # 创建React代理
                    self.agent = create_react_agent(model=llm, tools=self.tools, debug=True)
                    self.initialized = True
                    
                    # 恢复原始路径
                    sys.path = original_path
                    return self
        except Exception as e:
            print(f"初始化真实浏览器客户端失败: {e}")
            # 恢复原始路径
            sys.path = original_path
            # 如果真实浏览器初始化失败，回退到模拟版本
            return await self._fallback_initialize()
    
    async def _fallback_initialize(self):
        """回退到模拟浏览器客户端"""
        print("回退到模拟浏览器客户端")
        self.initialized = True
        return self
    
    async def browse_website(self, url: str, action: str = "获取页面内容") -> Dict[str, Any]:
        """浏览网站并执行操作"""
        if not self.initialized:
            await self.initialize()
        
        # 如果agent存在，使用真实浏览器功能
        if self.agent:
            try:
                prompt = f"请访问 {url} 并{action}"
                response = await self.agent.ainvoke(input={"messages": [("user", prompt)]})
                return self._parse_response(response)
            except Exception as e:
                print(f"真实浏览器操作失败: {e}")
                # 回退到模拟版本
                return await self._simulate_browse_website(url, action)
        else:
            # 使用模拟版本
            return await self._simulate_browse_website(url, action)
    
    async def search_and_extract(self, search_query: str, website: str = "百度", extract_info: str = "相关信息") -> Dict[str, Any]:
        """搜索并提取信息"""
        if not self.initialized:
            await self.initialize()
        
        # 如果agent存在，使用真实浏览器功能
        if self.agent:
            try:
                prompt = f"请在{website}中搜索'{search_query}'，并提取{extract_info}"
                response = await self.agent.ainvoke(input={"messages": [("user", prompt)]})
                return self._parse_response(response)
            except Exception as e:
                print(f"真实搜索操作失败: {e}")
                # 回退到模拟版本
                return await self._simulate_search_and_extract(search_query, website, extract_info)
        else:
            # 使用模拟版本
            return await self._simulate_search_and_extract(search_query, website, extract_info)
    
    async def _simulate_browse_website(self, url: str, action: str) -> Dict[str, Any]:
        """模拟浏览网站操作"""
        return {
            "success": True,
            "final_answer": f"已访问 {url} 并{action}",
            "tool_calls": [{"name": "browse_website", "args": {"url": url, "action": action}}]
        }
    
    async def _simulate_search_and_extract(self, search_query: str, website: str, extract_info: str) -> Dict[str, Any]:
        """模拟搜索并提取信息操作"""
        return {
            "success": True,
            "final_answer": f"已在{website}中搜索'{search_query}'，提取{extract_info}",
            "tool_calls": [{"name": "search_and_extract", "args": {"query": search_query, "website": website, "info": extract_info}}]
        }
    
    def _parse_response(self, response: Dict[str, Any]) -> Dict[str, Any]:
        """解析代理响应"""
        result = {
            "messages": [],
            "tool_calls": [],
            "final_answer": ""
        }
        
        # 简化解析逻辑，专注于获取最终答案
        if "messages" in response:
            for message in response["messages"]:
                if hasattr(message, 'content') and message.content:
                    result["final_answer"] = message.content
                    break
        
        return result


class SimpleSSEClient:
    """简化版SSE客户端，避免导入冲突"""
    
    def __init__(self):
        self.initialized = False
    
    async def initialize(self):
        """初始化SSE客户端"""
        self.initialized = True
        return self
    
    async def process_request(self, user_input: str) -> Dict[str, Any]:
        """处理SSE请求"""
        if not self.initialized:
            await self.initialize()
        
        return {
            "success": True,
            "final_answer": f"已处理SSE请求: {user_input}",
            "tool_calls": [{"name": "sse_process", "args": {"input": user_input}}]
        }


class SimpleAMAPClient:
    """简化版高德地图客户端，避免导入冲突"""
    
    def __init__(self):
        self.initialized = False
        self.amap_key = os.environ.get("AMAP_KEY")
    
    async def initialize(self):
        """初始化高德地图客户端"""
        self.initialized = True
        return self
    
    async def process_request(self, user_input: str) -> Dict[str, Any]:
        """处理高德地图请求"""
        if not self.initialized:
            await self.initialize()
        
        if not self.amap_key:
            return {"error": "未设置AMAP_KEY环境变量，高德地图功能不可用"}
        
        return {
            "success": True,
            "final_answer": f"已处理高德地图请求: {user_input}",
            "tool_calls": [{"name": "amap_process", "args": {"input": user_input, "key": self.amap_key}}]
        }


class ProjectMCPManager:
    """项目MCP管理器类，管理项目中已实现的MCP客户端"""
    
    def __init__(self):
        self.amap_client = None
        self.sse_client = None
        self.browser_client = None
        self.active_client = None
    
    async def initialize_amap_client(self):
        """初始化高德地图MCP客户端"""
        try:
            # 使用简化版客户端避免导入冲突
            self.amap_client = SimpleAMAPClient()
            await self.amap_client.initialize()
            return True
        except Exception as e:
            print(f"初始化高德地图MCP客户端失败: {e}")
            return False
    
    async def initialize_sse_client(self):
        """初始化SSE MCP客户端"""
        try:
            # 使用简化版客户端避免导入冲突
            self.sse_client = SimpleSSEClient()
            await self.sse_client.initialize()
            return True
        except Exception as e:
            print(f"初始化SSE MCP客户端失败: {e}")
            return False
    
    async def initialize_browser_client(self):
        """初始化浏览器MCP客户端"""
        try:
            # 使用真实浏览器客户端
            self.browser_client = RealBrowserClient()
            await self.browser_client.initialize()
            return True
        except Exception as e:
            print(f"初始化浏览器MCP客户端失败: {e}")
            return False
    
    async def process_amap_request(self, user_input: str) -> Dict[str, Any]:
        """处理高德地图相关请求"""
        if not self.amap_client:
            return {"error": "高德地图客户端未初始化"}
        
        try:
            result = await self.amap_client.process_request(user_input)
            return result
        except Exception as e:
            return {"error": f"处理高德地图请求时发生错误: {str(e)}"}
    
    async def process_sse_request(self, user_input: str) -> Dict[str, Any]:
        """处理SSE相关请求"""
        if not self.sse_client:
            return {"error": "SSE客户端未初始化"}
        
        try:
            result = await self.sse_client.process_request(user_input)
            return result
        except Exception as e:
            return {"error": f"处理SSE请求时发生错误: {str(e)}"}
    
    async def process_browser_request(self, user_input: str) -> Dict[str, Any]:
        """处理浏览器相关请求"""
        if not self.browser_client:
            return {"error": "浏览器客户端未初始化"}
        
        try:
            # 根据用户输入选择不同的浏览器操作
            if "搜索" in user_input or "查询" in user_input:
                # 提取搜索关键词
                if "天气" in user_input:
                    result = await self.browser_client.search_and_extract(
                        "天气", "百度", "天气信息"
                    )
                else:
                    # 默认搜索
                    result = await self.browser_client.search_and_extract(
                        user_input, "百度", "相关信息"
                    )
            elif "浏览" in user_input or "访问" in user_input:
                # 提取URL
                if "百度" in user_input:
                    result = await self.browser_client.browse_website(
                        "https://www.baidu.com", "获取页面内容"
                    )
                elif "新闻" in user_input:
                    result = await self.browser_client.browse_website(
                        "https://news.baidu.com", "获取新闻内容"
                    )
                else:
                    result = await self.browser_client.browse_website(
                        "https://www.baidu.com", "获取页面内容"
                    )
            else:
                # 默认操作
                result = await self.browser_client.browse_website(
                    "https://www.baidu.com", "获取页面内容"
                )
            
            return result
        except Exception as e:
            return {"error": f"处理浏览器请求时发生错误: {str(e)}"}


class AppStartup:
    """应用启动类，集成项目中已实现的MCP功能和AI对话"""
    
    def __init__(self, model: Optional[str] = None):
        self.model = model or settings.bailian.model
        # 设置百炼API密钥
        dashscope.api_key = settings.bailian.api_key
        
        # 检查API密钥是否设置
        if not dashscope.api_key:
            raise ValueError("未设置百炼API密钥。请检查.env文件中的BAILIAN_API_KEY设置。")
        
        # 初始化项目MCP管理器
        self.mcp_manager = ProjectMCPManager()
        self.conversation_history = []
    
    async def initialize_mcp_clients(self):
        """初始化项目中已实现的MCP客户端"""
        print("正在初始化项目中已实现的MCP客户端...")
        
        # 初始化高德地图客户端
        amap_success = await self.mcp_manager.initialize_amap_client()
        if amap_success:
            print("✓ 高德地图MCP客户端初始化成功")
        else:
            print("✗ 高德地图MCP客户端初始化失败（需要设置AMAP_KEY）")
        
        # 初始化SSE客户端
        sse_success = await self.mcp_manager.initialize_sse_client()
        if sse_success:
            print("✓ SSE MCP客户端初始化成功")
        else:
            print("✗ SSE MCP客户端初始化失败")
        
        # 初始化浏览器客户端
        browser_success = await self.mcp_manager.initialize_browser_client()
        if browser_success:
            print("✓ 浏览器MCP客户端初始化成功")
        else:
            print("✗ 浏览器MCP客户端初始化失败")
        
        return amap_success or sse_success or browser_success
    
    async def chat_with_ai(self, user_message: str) -> str:
        """使用百炼平台进行AI对话"""
        try:
            response = Generation.call(
                model=self.model,
                messages=[
                    {"role": "system", "content": "你是一个有用的助手，可以协助用户进行地图查询、浏览器操作等任务。"},
                    *self.conversation_history,
                    {"role": "user", "content": user_message},
                ],
                result_format='message'
            )
            
            if response.status_code == 200:
                ai_response = response.output.choices[0].message.content
                # 更新对话历史
                self.conversation_history.append({"role": "user", "content": user_message})
                self.conversation_history.append({"role": "assistant", "content": ai_response})
                return ai_response
            else:
                return f"请求失败，状态码: {response.status_code}, 错误信息: {response.message}"
                
        except Exception as e:
            return f"发生错误: {str(e)}"
    
    async def process_user_input(self, user_input: str) -> Dict[str, Any]:
        """处理用户输入，智能选择使用AI对话或项目MCP功能"""
        # 检查是否需要使用项目MCP功能
        amap_keywords = ["地图", "导航", "位置", "路线", "高德", "amap"]
        sse_keywords = ["计算", "数学", "sse", "stream"]
        browser_keywords = ["浏览", "搜索", "网站", "网页", "浏览器"]
        
        if any(keyword in user_input for keyword in amap_keywords):
            # 使用高德地图MCP功能
            result = await self.mcp_manager.process_amap_request(user_input)
        elif any(keyword in user_input for keyword in sse_keywords):
            # 使用SSE MCP功能
            result = await self.mcp_manager.process_sse_request(user_input)
        elif any(keyword in user_input for keyword in browser_keywords):
            # 使用浏览器MCP功能
            result = await self.mcp_manager.process_browser_request(user_input)
        else:
            # 使用AI对话
            ai_response = await self.chat_with_ai(user_input)
            result = {
                "success": True,
                "final_answer": ai_response,
                "type": "ai"
            }
        
        # 将结果添加到对话历史（如果不是错误）
        if "error" not in result:
            self.conversation_history.append({"role": "user", "content": user_input})
            if result.get("final_answer"):
                self.conversation_history.append({"role": "assistant", "content": result["final_answer"]})
        
        return result


async def interactive_chat():
    """交互式多轮对话"""
    app = AppStartup()
    
    # 初始化项目中已实现的MCP客户端
    mcp_initialized = await app.initialize_mcp_clients()
    
    print("\n" + "="*60)
    print("AI助手 + 项目MCP功能集成系统")
    print("="*60)
    print("系统已启动！您可以：")
    print("1. 直接与我对话")
    print("2. 使用高德地图功能（输入包含'地图'、'导航'、'位置'等关键词）")
    print("3. 使用SSE功能（输入包含'计算'、'数学'等关键词）")
    print("4. 使用浏览器功能（输入包含'浏览'、'搜索'、'网站'等关键词）")
    print("输入'退出'或'quit'结束对话")
    print("="*60)
    
    while True:
        try:
            user_input = input("\n您: ").strip()
            
            if user_input.lower() in ['退出', 'quit', 'exit']:
                print("再见！")
                break
            
            if not user_input:
                continue
            
            print("系统: ", end="", flush=True)
            
            # 处理用户输入
            result = await app.process_user_input(user_input)
            
            if "error" in result:
                print(f"错误: {result['error']}")
            elif result.get("success", False):
                if result.get("final_answer"):
                    print(result["final_answer"])
                else:
                    print("操作已完成")
                    
                # 显示工具调用信息
                if result.get("tool_calls"):
                    print("\n[工具调用记录]")
                    for tool_call in result["tool_calls"]:
                        print(f"- 调用工具: {tool_call['name']}")
            
        except KeyboardInterrupt:
            print("\n\n再见！")
            break
        except Exception as e:
            print(f"\n发生错误: {e}")


async def main():
    """主函数"""
    await interactive_chat()


if __name__ == "__main__":
    asyncio.run(main())